package com.example.edu.Utils;

import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.embedding.Embedding;
import com.zhipu.oapi.service.v4.embedding.EmbeddingApiResponse;
import com.zhipu.oapi.service.v4.embedding.EmbeddingRequest;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

import static com.example.edu.Utils.MyConstants.EMBEDDING_API_KEY;

@Component
public class EmbeddingUtil {
    public List<List<Double>> getEmbeddings(List<String> texts){
        ClientV4 client = new ClientV4.Builder(EMBEDDING_API_KEY)
                .networkConfig(300,30,160,110, TimeUnit.SECONDS)
                .build();
        EmbeddingRequest embeddingRequest = new EmbeddingRequest();
        embeddingRequest.setInput(texts);
        embeddingRequest.setModel(Constants.ModelEmbedding3);
        EmbeddingApiResponse apiResponse = client.invokeEmbeddingsApi(embeddingRequest);
        List<Embedding> embeddings = apiResponse.getData().getData();
        List<List<Double>> r = new ArrayList<>();
        for (Embedding embedding : embeddings)
            r.add(embedding.getEmbedding());
        return r;
    }
}
